import torch
from torch import nn
import torch.nn.functional as F


# class Attention(nn.Module):
#     def __init__(self, dim, head, sr_ratio):
#         """
#         注意力头
#         :param dim: 输入维度
#         :param head: 注意力头数目
#         :param sr_ratio: 缩放倍数
#         """
#         super().__init__()
#         self.head = head
#         self.sr_ratio = sr_ratio
#         self.scale = (dim // head) ** -0.5
#         self.q = nn.Linear(dim, dim)
#         self.kv = nn.Linear(dim, dim * 2)
#         self.proj = nn.Linear(dim, dim)

#         if sr_ratio > 1:
#             self.sr = nn.Conv2d(dim, dim, sr_ratio, sr_ratio)
#             self.norm = nn.LayerNorm(dim)

#     def forward(self, x, y, H, W):

#         B,C,H,W 
#         B, N, C = x.shape
#         q = self.q(x).reshape(B, N, self.head, C // self.head).permute(0, 2, 1, 3)

#         if self.sr_ratio > 1:
#             x = x.permute(0, 2, 1).reshape(B, C, H, W)
#             x = self.sr(x).reshape(B, C, -1).permute(0, 2, 1)
#             x = self.norm(x)

#         k, v = self.kv(x).reshape(B, -1, 2, self.head, C // self.head).permute(2, 0, 3, 1, 4)

#         attn = (q @ k.transpose(-2, -1)) * self.scale
#         attn = attn.softmax(dim=-1)

#         x = (attn @ v).transpose(1, 2).reshape(B, N, C)
#         x = self.proj(x)
#         return x

# class Attention(nn.Module):
#     def __init__(self, dim, num_heads=8, head_dim_ratio=1., attn_drop=0., proj_drop=0.):
#         super().__init__()
#         self.dim = dim
#         self.num_heads = num_heads
#         head_dim = round(dim // num_heads * head_dim_ratio)
#         self.head_dim = head_dim
#         self.scale = head_dim ** -0.5
#         self.qkv = nn.Conv2d(dim, head_dim * num_heads * 3, 1, stride=1, padding=0, bias=False)
#         self.attn_drop = nn.Dropout(attn_drop)
#         self.proj = nn.Conv2d(self.head_dim * self.num_heads, dim, 1, stride=1, padding=0, bias=False)
#         self.proj_drop = nn.Dropout(proj_drop)

#     def forward(self, x):
#         B, C, H, W = x.shape
#         x = self.qkv(x).reshape(B, 3, self.num_heads, self.head_dim, -1).permute(1, 0, 2, 4, 3) # [3,4,8,3600,64]
#         q, k, v = x.unbind(0)

#         attn = (q @ k.transpose(-2, -1)) * self.scale
#         attn = attn.softmax(dim=-1)
#         attn = self.attn_drop(attn)
#         x = attn @ v

#         x = x.permute(0, 1, 3, 2).reshape(B, -1, H, W)
#         x = self.proj(x)
#         x = self.proj_drop(x)
#         return x
    
class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, head_dim_ratio=1., attn_drop=0., proj_drop=0.):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        head_dim = round(dim // num_heads * head_dim_ratio)
        self.head_dim = head_dim
        self.scale = head_dim ** -0.5
        self.qk = nn.Conv2d(dim, head_dim * num_heads * 2, 1, stride=1, padding=0, bias=False)
        self.v = nn.Conv2d(dim, head_dim * num_heads, 1, stride=1, padding=0, bias=False)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Conv2d(self.head_dim * self.num_heads, dim, 1, stride=1, padding=0, bias=False)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x , query):
        B, C, H, W = x.shape
        x = self.qk(x).reshape(B, 2, self.num_heads, self.head_dim, -1).permute(1, 0, 2, 4, 3)
        q, k = x.unbind(0)

        v = self.v(query).reshape(B, self.num_heads, self.head_dim, -1).permute(0, 1, 3, 2)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = attn @ v

        x = x.permute(0, 1, 3, 2).reshape(B, -1, H, W)
        x = self.proj(x)
        x = self.proj_drop(x) + query
        return x
    
if __name__=="__main__":

    x=torch.randn(4,512,60,60)
    query= torch.randn(4,512,60,60)
    model=Attention(512)
    y=model(x,query)
    print(y.shape)

    